package com.it.service.impl;

import com.it.pojo.vo.KnowledgeFileVO;
import com.it.service.KnowledgeService;
import com.it.service.RagService;
import com.it.utils.RagFileUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

/**
 * RAG服务实现类
 * 负责处理知识库文件的向量化存储和检索
 */
@Service
@Slf4j
public class RagServiceImpl implements RagService {

    @Autowired
    private VectorStore vectorStore;

    @Autowired
    private KnowledgeService knowledgeService;

    @Autowired
    private RagFileUtil ragFileUtil;

    // 用于跟踪已构建向量数据库的知识库
    private final Set<String> builtKnowledgeBases = ConcurrentHashMap.newKeySet();

    /**
     * 根据知识库编码构建向量数据库
     * 将指定知识库下的所有文件内容向量化并存储到向量数据库中
     *
     * @param knowledgeCode 知识库编码
     * @return 成功处理的文件数量
     */
    @Override
    public int buildVectorStore(String knowledgeCode) {
        log.info("开始为知识库 {} 构建向量数据库", knowledgeCode);

        try {
            // 获取知识库下的所有文件
            List<KnowledgeFileVO> fileList = knowledgeService.getKnowledgeFilesByCode(knowledgeCode);
            
            if (fileList == null || fileList.isEmpty()) {
                log.warn("知识库 {} 下没有文件，无需构建向量数据库", knowledgeCode);
                return 0;
            }

            List<Document> documents = new ArrayList<>();
            int successCount = 0;

            for (KnowledgeFileVO file : fileList) {
                try {
                    log.info("处理文件: {}", file.getFileName());
                    
                    // 读取文件内容
                    String content = ragFileUtil.getFileByPath(file.getFilePath());
                    
                    if (content == null || content.trim().isEmpty()) {
                        log.warn("文件 {} 内容为空，跳过处理", file.getFileName());
                        continue;
                    }

                    // 创建文档对象，包含文件内容和元数据
                    Map<String, Object> metadata = new HashMap<>();
                    metadata.put("knowledgeCode", knowledgeCode);
                    metadata.put("fileName", file.getFileName());
                    metadata.put("filePath", file.getFilePath());
                    metadata.put("fileType", file.getFileType());
                    metadata.put("fileId", file.getId().toString());

                    // 将长文本分块处理（每块最大2000字符）
                    List<String> chunks = splitTextIntoChunks(content, 2000);
                    
                    for (int i = 0; i < chunks.size(); i++) {
                        Map<String, Object> chunkMetadata = new HashMap<>(metadata);
                        chunkMetadata.put("chunkIndex", i);
                        chunkMetadata.put("totalChunks", chunks.size());
                        
                        Document document = new Document(chunks.get(i), chunkMetadata);
                        documents.add(document);
                    }

                    successCount++;
                    log.info("文件 {} 处理完成，分为 {} 个文本块", file.getFileName(), chunks.size());

                } catch (Exception e) {
                    log.error("处理文件 {} 时发生错误: {}", file.getFileName(), e.getMessage(), e);
                }
            }

            if (!documents.isEmpty()) {
                // 将文档添加到向量数据库
                vectorStore.add(documents);
                log.info("成功将 {} 个文档块添加到向量数据库", documents.size());
                
                // 标记该知识库已构建向量数据库
                builtKnowledgeBases.add(knowledgeCode);
            }

            log.info("知识库 {} 向量数据库构建完成，成功处理 {} 个文件", knowledgeCode, successCount);
            return successCount;

        } catch (Exception e) {
            log.error("构建知识库 {} 的向量数据库时发生错误: {}", knowledgeCode, e.getMessage(), e);
            throw new RuntimeException("构建向量数据库失败: " + e.getMessage(), e);
        }
    }

    /**
     * 根据查询内容在向量数据库中检索相关文档
     *
     * @param query 查询内容
     * @param knowledgeCode 知识库编码（用于过滤特定知识库的内容）
     * @param topK 返回最相关的文档数量
     * @return 相关文档内容列表
     */
    @Override
    public List<String> searchSimilarDocuments(String query, String knowledgeCode, int topK) {
        log.info("在知识库 {} 中搜索相关文档，查询内容: {}, topK: {}", knowledgeCode, query, topK);

        try {
            // 创建搜索请求
            SearchRequest searchRequest = SearchRequest.builder()
                    .query(query)
                    .build();
            
            // 执行相似性搜索
            List<Document> similarDocuments = vectorStore.similaritySearch(searchRequest);
            
            // 过滤出属于指定知识库的文档，并提取内容
            List<String> results = similarDocuments.stream()
                    .filter(doc -> {
                        Object docKnowledgeCode = doc.getMetadata().get("knowledgeCode");
                        return docKnowledgeCode != null && docKnowledgeCode.toString().equals(knowledgeCode);
                    })
                    .map(Document::getFormattedContent)
                    .collect(Collectors.toList());

            log.info("找到 {} 个相关文档片段", results.size());
            return results;

        } catch (Exception e) {
            log.error("搜索相关文档时发生错误: {}", e.getMessage(), e);
            return new ArrayList<>();
        }
    }

    /**
     * 清空指定知识库的向量数据
     *
     * @param knowledgeCode 知识库编码
     */
    @Override
    public void clearVectorStore(String knowledgeCode) {
        log.info("清空知识库 {} 的向量数据", knowledgeCode);
        
        try {
            // 注意：SimpleVectorStore 没有直接的删除方法
            // 这里我们只是从已构建列表中移除，实际的向量数据会在重新构建时覆盖
            builtKnowledgeBases.remove(knowledgeCode);
            log.info("已从已构建列表中移除知识库 {}", knowledgeCode);
            
        } catch (Exception e) {
            log.error("清空知识库 {} 向量数据时发生错误: {}", knowledgeCode, e.getMessage(), e);
        }
    }

    /**
     * 检查指定知识库是否已经构建了向量数据库
     *
     * @param knowledgeCode 知识库编码
     * @return 是否已构建
     */
    @Override
    public boolean isVectorStoreBuilt(String knowledgeCode) {
        return builtKnowledgeBases.contains(knowledgeCode);
    }

    /**
     * 将长文本分割成较小的块
     *
     * @param text 原始文本
     * @param maxChunkSize 每块的最大字符数
     * @return 文本块列表
     */
    private List<String> splitTextIntoChunks(String text, int maxChunkSize) {
        List<String> chunks = new ArrayList<>();
        
        if (text.length() <= maxChunkSize) {
            chunks.add(text);
            return chunks;
        }

        // 按段落分割
        String[] paragraphs = text.split("\n\n");
        StringBuilder currentChunk = new StringBuilder();

        for (String paragraph : paragraphs) {
            // 如果当前块加上新段落超过限制，则保存当前块并开始新块
            if (currentChunk.length() + paragraph.length() > maxChunkSize && currentChunk.length() > 0) {
                chunks.add(currentChunk.toString().trim());
                currentChunk = new StringBuilder();
            }

            // 如果单个段落就超过限制，需要进一步分割
            if (paragraph.length() > maxChunkSize) {
                // 保存当前块（如果有内容）
                if (currentChunk.length() > 0) {
                    chunks.add(currentChunk.toString().trim());
                    currentChunk = new StringBuilder();
                }
                
                // 按句子分割长段落
                String[] sentences = paragraph.split("[。！？.!?]");
                for (String sentence : sentences) {
                    if (currentChunk.length() + sentence.length() > maxChunkSize && currentChunk.length() > 0) {
                        chunks.add(currentChunk.toString().trim());
                        currentChunk = new StringBuilder();
                    }
                    currentChunk.append(sentence).append("。");
                }
            } else {
                currentChunk.append(paragraph).append("\n\n");
            }
        }

        // 添加最后一块
        if (currentChunk.length() > 0) {
            chunks.add(currentChunk.toString().trim());
        }

        return chunks;
    }
}